# Dickstein, Ho, and Mark (2023)

is.inf <- function(x) (x %in% c(-Inf, Inf))

parametrize_cost_data <- function(
  cost_theta,
  data,
  cost_covar_names,
  cost_covar_lens,
  fixed_cost_censor=NULL,
  sim=FALSE) {
  
  if (max(grepl("^score", names(data)))){
    stop("There is a column name starting with 'score' here which is not ok. Please amend your ways.")
  }

  data <- data %>% filter(choice == 1)
  
  # Back out individual covar names and beta params from covar_lanes
  W1_names <- cost_covar_names[1:cost_covar_lens[1]]
  W2_names <- cost_covar_names[(1+cost_covar_lens[1]):(cost_covar_lens[1]+cost_covar_lens[2])]
  
  beta1 <- cost_theta[1:cost_covar_lens[1]]
  beta2 <- cost_theta[(1+cost_covar_lens[1]):(cost_covar_lens[1]+cost_covar_lens[2])]
  
  
  if (is.null(fixed_cost_censor)) {
    # If trying to estimate cost threshold, take last param as cost threshold
    cost_censor <- cost_theta[length(cost_theta)] 
    cost_censor <- max(cost_censor, 0)
  } else {
    cost_censor <- fixed_cost_censor
  }
  
  # Construct covar vectors
  W1_mat <- as.matrix(data[W1_names])
  W2_mat <- as.matrix(data[W2_names])
  
  # Construct parameter vectors
  data['alpha_i'] <- exp(W1_mat %*% beta1)
  data['omega_i'] <- exp(W2_mat %*% beta2)
  
  # Construct conditional-lambda==1 cost 
  data['lambda_one_cost'] <- data['x_av'] + (data['omega_i'] * data['x_av']^2)

  # Construct "probability of cost being below the censor"
  data['pr_high_cost'] <- exp(- cost_censor * data['alpha_i']/data['lambda_one_cost']) 
  data['pr_low_cost'] <- 1 - data['pr_high_cost']

  # Construct individual-LLH conditional on spending
  data <- data %>%
    mutate(
      healthy_ind = as.numeric(cost <= cost_censor),
      log_f_c_ij_sick = log(alpha_i/lambda_one_cost) - (cost * (alpha_i/lambda_one_cost)),
      log_f_c_ij = case_when(
        (x_av > 0 & healthy_ind == 1) ~ log(pr_low_cost),
        (x_av > 0 & healthy_ind == 0) ~ log_f_c_ij_sick))
  if(any(is.inf(data$log_f_c_ij))){
    warning("There are cost probabilities that are equal to zero in the parametrized cost data.")
  }

  # Construct gradient
  # beta_1 derivatives
  for (i in 1:cost_covar_lens[1]) {
    healthy_score_W1 <- data['alpha_i'] * data[W1_names[i]] *
      (data['pr_high_cost']/data['pr_low_cost']) * 
      cost_censor / data['lambda_one_cost']
    sick_score_W1 <- data[W1_names[i]] * (
      1 - (data['alpha_i']*data['cost']/data['lambda_one_cost']))
    data[paste('cost_score_W1', i, sep='_')] <- data['choice'] * (
      data['healthy_ind']*healthy_score_W1 + (1-data['healthy_ind'])*sick_score_W1)
  }
    
  # beta_2 derivatives
  for (i in 1:cost_covar_lens[2]) {
    healthy_score_W2 <- - data['alpha_i'] * cost_censor * data['omega_i'] * data[W2_names[i]] *
      (data['pr_high_cost']/data['pr_low_cost']) *
      (data['x_av']/data['lambda_one_cost']) ^ 2
    sick_score_W2 <- (data['omega_i']*data[W2_names[i]] * data['x_av']^2) * (
      - (1/data['lambda_one_cost']) 
      + (data['cost']*data['alpha_i'] /(data['lambda_one_cost'] ^ 2)))
     data[paste('cost_score_W2', i, sep='_')] <- data['choice'] * ( 
       data['healthy_ind']*healthy_score_W2 + (1-data['healthy_ind'])*sick_score_W2) 
  }
  return(data)
}

